# qa_web.py
import os
import uuid
from flask import Flask, request, render_template_string, jsonify, session
# 导入增强版 HistoryMemory
from history1 import HistoryMemory
# --- 导入你的 Agent 系统 ---
# 假设 agent2.py 在同一个目录下，或者在 Python 路径中
from agent2 import EnterpriseAgentSystem  # 根据你的实际文件名调整

app = Flask(__name__)
# 设置一个密钥以启用 Flask sessions (用于存储当前会话 ID)
app.secret_key = os.environ.get("SECRET_KEY", "your_default_secret_key_change_this")

# --- 初始化 HistoryMemory ---
# 可以根据需要调整参数
history_memory = HistoryMemory(max_sessions=10, max_turns_per_session=20)

# --- 初始化 EnterpriseAgentSystem ---
# 请根据你的实际 FAISS 索引路径进行修改
FAISS_INDEX_PATH = "faiss_index"  # <--- 修改为你的实际路径
try:
    enterprise_agent = EnterpriseAgentSystem(faiss_index_path=FAISS_INDEX_PATH)
    print("✅ EnterpriseAgentSystem 初始化成功")
except Exception as e:
    print(f"❌ EnterpriseAgentSystem 初始化失败: {e}")
    enterprise_agent = None  # 如果初始化失败，后续调用需要检查


# --- 真实的 Agent 执行函数 ---
def real_agent_execute(query, context=""):
    """
    使用真实的 EnterpriseAgentSystem 执行查询，并利用上下文。
    """
    if not enterprise_agent:
        return "❌ Agent 系统未初始化，请检查日志。"

    # 1. 构造包含上下文的完整提示
    #    这里我们简单地将上下文作为前缀添加到查询中。
    #    你的 Ollama 模型 (qwen3:4b) 应该能够处理这种格式的长输入。
    full_query = f"请基于以下对话历史和文档知识回答问题。\n\n对话历史:\n{context}\n\n问题:\n{query}"

    # 2. 调用你的 Agent
    #    假设 execute_query 可以处理这种格式的查询。
    #    如果你的 Agent 内部已经处理了文档检索和 Prompt 构造，
    #    你可以直接传入 full_query。
    #    或者，如果你需要更精细的控制，可以在这里调用 DocumentStorage
    #    和 ReactAgentSystem 的方法。
    try:
        # 这里我们假设 execute_query 会处理所有逻辑
        # 如果需要区分 React 和 ZeroShot，可以根据 context 或其他逻辑决定
        # 注意：execute_query 会打印到 stdout，这里我们捕获返回值
        agent_response = enterprise_agent.execute_query(full_query, user_react=True)
        return agent_response
    except Exception as e:
        print(f"Agent 执行出错: {e}")
        return f"❌ 执行查询时发生错误: {str(e)}"


# --- HTML 模板 ---
HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>企业问答系统</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            margin: 0;
            padding: 0;
            display: flex;
            flex-direction: column;
            height: 100vh;
            background-color: #f5f5f5;
        }
        .header {
            background-color: #4CAF50;
            color: white;
            padding: 15px;
            text-align: center;
            font-size: 1.5em;
        }
        .chat-container {
            flex: 1;
            overflow-y: auto;
            padding: 20px;
            display: flex;
            flex-direction: column;
        }
        .message {
            max-width: 70%;
            padding: 10px 15px;
            margin-bottom: 15px;
            border-radius: 15px;
            position: relative;
            word-wrap: break-word;
            animation: fadeIn 0.3s ease-in-out;
        }
        @keyframes fadeIn {
            from { opacity: 0; transform: translateY(10px); }
            to { opacity: 1; transform: translateY(0); }
        }
        .user-message {
            background-color: #4CAF50;
            color: white;
            align-self: flex-end;
            border-bottom-right-radius: 5px;
        }
        .agent-message {
            background-color: #E0E0E0;
            color: black;
            align-self: flex-start;
            border-bottom-left-radius: 5px;
        }
        .message-info {
            font-size: 0.75em;
            color: #777;
            margin-top: 5px;
            text-align: right;
        }
        .input-container {
            display: flex;
            padding: 15px;
            background-color: white;
            border-top: 1px solid #ddd;
        }
        #user-input {
            flex: 1;
            padding: 10px;
            border: 1px solid #ccc;
            border-radius: 20px;
            outline: none;
        }
        #send-button {
            margin-left: 10px;
            padding: 10px 20px;
            background-color: #4CAF50;
            color: white;
            border: none;
            border-radius: 20px;
            cursor: pointer;
            outline: none;
        }
        #send-button:hover {
            background-color: #45a049;
        }
        #send-button:disabled {
            background-color: #cccccc;
            cursor: not-allowed;
        }
        .session-selector {
            padding: 10px;
            background-color: white;
            border-bottom: 1px solid #ddd;
            display: flex;
            align-items: center;
        }
        #session-select {
            flex: 1;
            padding: 5px;
            margin-right: 10px;
        }
        #new-session-button {
            padding: 5px 10px;
            background-color: #2196F3;
            color: white;
            border: none;
            border-radius: 5px;
            cursor: pointer;
        }
        #new-session-button:hover {
            background-color: #1976D2;
        }
        .typing-indicator {
            display: none;
            align-self: flex-start;
            background-color: #E0E0E0;
            color: black;
            padding: 10px 15px;
            margin-bottom: 15px;
            border-radius: 15px;
            border-bottom-left-radius: 5px;
        }
        .typing-dots {
            display: inline-block;
        }
        .typing-dots span {
            animation: typing 1.4s infinite ease-in-out;
            background-color: #666;
            border-radius: 50%;
            display: inline-block;
            height: 6px;
            width: 6px;
            margin: 0 2px;
        }
        .typing-dots span:nth-child(1) { animation-delay: 0s; }
        .typing-dots span:nth-child(2) { animation-delay: 0.2s; }
        .typing-dots span:nth-child(3) { animation-delay: 0.4s; }

        @keyframes typing {
            0%, 60%, 100% { transform: translateY(0); }
            30% { transform: translateY(-5px); }
        }
    </style>
</head>
<body>

<div class="header">
    企业问答系统
</div>

<div class="session-selector">
    <select id="session-select">
        <!-- 会话选项将通过 JS 动态加载 -->
    </select>
    <button id="new-session-button">新建会话</button>
</div>

<div class="chat-container" id="chat-container">
    <div class="typing-indicator" id="typing-indicator">
        <div class="typing-dots">
            <span></span>
            <span></span>
            <span></span>
        </div>
    </div>
    <!-- 消息将通过 JS 动态添加 -->
</div>

<div class="input-container">
    <input type="text" id="user-input" placeholder="输入您的问题..." autocomplete="off">
    <button id="send-button">发送</button>
</div>

<script>
    let currentSessionId = null;

    // --- API 调用函数 ---
    async function fetchSessions() {
        const response = await fetch('/api/sessions');
        return await response.json();
    }

    async function createSession() {
        const response = await fetch('/api/sessions', { method: 'POST' });
        return await response.json();
    }

    async function sendMessage(message, sessionId) {
        const response = await fetch('/api/chat', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({ message: message, session_id: sessionId })
        });
        if (!response.ok) {
            throw new Error(`HTTP error! status: ${response.status}`);
        }
        return await response.json();
    }

    // --- UI 更新函数 ---
    function updateSessionSelector(sessions) {
        const selectElement = document.getElementById('session-select');
        selectElement.innerHTML = '';
        sessions.forEach(session => {
            const option = document.createElement('option');
            option.value = session.id;
            // 简化显示，只显示 ID 的一部分
            option.textContent = `会话 ${session.id.substring(0, 8)}`;
            selectElement.appendChild(option);
        });
        if (currentSessionId && sessions.some(s => s.id === currentSessionId)) {
            selectElement.value = currentSessionId;
        } else if (sessions.length > 0) {
            currentSessionId = sessions[0].id;
            selectElement.value = currentSessionId;
        } else {
            currentSessionId = null;
        }
    }

    function addMessageToChat(message, sender, timestamp='') {
        const chatContainer = document.getElementById('chat-container');
        const messageDiv = document.createElement('div');
        messageDiv.className = `message ${sender}-message`;

        const contentDiv = document.createElement('div');
        contentDiv.textContent = message;
        messageDiv.appendChild(contentDiv);

        if (timestamp) {
            const infoDiv = document.createElement('div');
            infoDiv.className = 'message-info';
            infoDiv.textContent = timestamp;
            messageDiv.appendChild(infoDiv);
        }

        chatContainer.appendChild(messageDiv);
        chatContainer.scrollTop = chatContainer.scrollHeight;
    }

    function showTypingIndicator() {
        document.getElementById('typing-indicator').style.display = 'block';
        document.getElementById('chat-container').scrollTop = document.getElementById('chat-container').scrollHeight;
    }

    function hideTypingIndicator() {
        document.getElementById('typing-indicator').style.display = 'none';
    }

    // --- 事件处理函数 ---
    async function handleNewSession() {
        const data = await createSession();
        currentSessionId = data.session_id;
        await loadSessionsAndRefresh(); // 刷新会话列表并加载新会话的历史
    }

    async function handleSendMessage() {
        const inputElement = document.getElementById('user-input');
        const message = inputElement.value.trim();
        if (!message || !currentSessionId) return;

        // 添加用户消息到 UI
        const now = new Date();
        const timeString = `${now.getHours().toString().padStart(2, '0')}:${now.getMinutes().toString().padStart(2, '0')}:${now.getSeconds().toString().padStart(2, '0')}`;
        addMessageToChat(message, 'user', timeString);

        // 清空输入框并禁用按钮
        inputElement.value = '';
        const sendButton = document.getElementById('send-button');
        sendButton.disabled = true;
        showTypingIndicator();

        try {
            // 发送消息到后端
            const data = await sendMessage(message, currentSessionId);

            // 添加 Agent 回复到 UI
            const responseTime = data.timestamp || new Date().toISOString();
            const responseTimeString = new Date(responseTime).toLocaleTimeString('zh-CN', { hour12: false });
            addMessageToChat(data.response, 'agent', responseTimeString);
        } catch (error) {
            console.error('发送消息失败:', error);
            addMessageToChat(`❌ 发送消息失败: ${error.message}`, 'agent');
        } finally {
            // 重新启用按钮并隐藏加载指示器
            sendButton.disabled = false;
            hideTypingIndicator();
            inputElement.focus();
        }
    }

    async function loadSessionsAndRefresh() {
        try {
            const sessions = await fetchSessions();
            updateSessionSelector(sessions);
            await loadChatHistory(currentSessionId); // 加载选定会话的历史
        } catch (error) {
            console.error('加载会话失败:', error);
        }
    }

    async function loadChatHistory(sessionId) {
        if (!sessionId) {
             document.getElementById('chat-container').innerHTML = '<div class="typing-indicator" id="typing-indicator"><div class="typing-dots"><span></span><span></span><span></span></div></div>';
            return;
        }
        try {
            const response = await fetch(`/api/history/${sessionId}`);
            if (!response.ok) {
                 document.getElementById('chat-container').innerHTML = '<div class="typing-indicator" id="typing-indicator"><div class="typing-dots"><span></span><span></span><span></span></div></div>';
                throw new Error(`HTTP error! status: ${response.status}`);
            }
            const history = await response.json();
            const chatContainer = document.getElementById('chat-container');
            chatContainer.innerHTML = '<div class="typing-indicator" id="typing-indicator"><div class="typing-dots"><span></span><span></span><span></span></div></div>'; // 清空并保留加载指示器
            history.forEach(turn => {
                const userTime = turn.user.timestamp ? new Date(turn.user.timestamp).toLocaleTimeString('zh-CN', { hour12: false }) : '';
                const agentTime = turn.agent.timestamp ? new Date(turn.agent.timestamp).toLocaleTimeString('zh-CN', { hour12: false }) : '';
                addMessageToChat(turn.user.message, 'user', userTime);
                addMessageToChat(turn.agent.message, 'agent', agentTime);
            });
            chatContainer.scrollTop = chatContainer.scrollHeight;
        } catch (error) {
            console.error('加载聊天历史失败:', error);
             document.getElementById('chat-container').innerHTML = '<div class="typing-indicator" id="typing-indicator"><div class="typing-dots"><span></span><span></span><span></span></div></div>';
        }
    }

    // --- 页面加载和事件监听 ---
    document.addEventListener('DOMContentLoaded', async () => {
        await loadSessionsAndRefresh(); // 初始加载会话

        document.getElementById('new-session-button').addEventListener('click', handleNewSession);

        document.getElementById('session-select').addEventListener('change', async (event) => {
            currentSessionId = event.target.value;
            await loadChatHistory(currentSessionId);
        });

        document.getElementById('send-button').addEventListener('click', handleSendMessage);
        document.getElementById('user-input').addEventListener('keypress', (e) => {
            if (e.key === 'Enter') {
                handleSendMessage();
            }
        });
    });

</script>

</body>
</html>
"""


# --- Flask 路由 ---
@app.route("/")
def index():
    """主页: 渲染聊天界面"""
    return render_template_string(HTML_TEMPLATE)


# --- API 路由 ---

@app.route("/api/chat", methods=["POST"])
def api_chat():
    """API: 处理聊天消息"""
    data = request.get_json()
    user_message = data.get("message")
    session_id = data.get("session_id")

    if not user_message or not session_id:
        return jsonify({"error": "缺少消息或会话ID"}), 400

    # 1. 获取上下文
    context = history_memory.get_context(session_id, last_n=5)  # 获取最近5轮上下文，可根据需要调整

    # 2. 调用你的 Agent (使用真实函数)
    # agent_response = mock_agent_execute(user_message, context=context) # <-- 注释掉模拟函数
    agent_response = real_agent_execute(user_message, context=context)  # <-- 使用真实函数

    # 3. 将对话轮次添加到历史记录
    history_memory.add_turn(user_message, agent_response, session_id)

    # 4. 返回响应
    # 注意：实际应用中，Agent 可能会返回时间戳，这里我们用当前时间
    import datetime
    response_timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return jsonify({"response": agent_response, "timestamp": response_timestamp})


@app.route("/api/sessions", methods=["GET"])
def api_get_sessions():
    """API: 获取所有会话 ID"""
    sessions = history_memory.get_all_sessions()
    # 转换为前端期望的格式
    session_list = [{"id": sid} for sid in sessions]
    return jsonify(session_list)


@app.route("/api/sessions", methods=["POST"])
def api_create_session():
    """API: 创建新的会话 ID"""
    new_session_id = str(uuid.uuid4())
    # 可以选择性地在这里初始化会话，虽然 HistoryMemory 会在 add_turn 时自动处理
    return jsonify({"session_id": new_session_id})


@app.route("/api/history/<session_id>", methods=["GET"])
def api_get_history(session_id):
    """API: 获取特定会话的历史记录"""
    history = history_memory.get_history(session_id)
    if history is None:
        return jsonify({"error": "会话未找到"}), 404
    return jsonify(history)


if __name__ == "__main__":
    # 确保在非调试模式下也能重新加载模块（如果需要）
    # debug=True 会启用重载，但注意它会重启应用，可能影响 session 状态
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 5000)), debug=False)
